from Data.dataset_process import load_dataset
from sklearn.model_selection import KFold
from sklearn import preprocessing
from sklearn import neighbors
from sklearn.metrics import roc_curve
import matplotlib.pyplot as plt
import numpy as np


def sk_predictor():
    data_types = ["HepTh"]
    dataset = load_dataset(data_types[0])
    x = dataset.data
    y = dataset.target
    # y = [1 if value == 'True' else 0 for value in dataset.target]
    min_max = preprocessing.MinMaxScaler()
    x = min_max.fit_transform(x)
    k_fold = KFold(n_splits=10, shuffle=True)
    fprs = []
    tprs = []

    for (train, test) in k_fold.split(x):
        x_train = [x[index] for index in train]
        x_test = [x[index] for index in test]
        y_train = [y[index] for index in train]
        y_test = [y[index] for index in test]

        model = neighbors.KNeighborsClassifier(n_neighbors=10, n_jobs=10)
        model.fit(x_train, y_train)
        y_score = model.predict_proba(x_test)
        fpr, tpr, _ = roc_curve(y_test, y_score[:, 1])
        fpr = fpr.tolist()
        tpr = tpr.tolist()
        fprs.append(fpr)
        tprs.append(tpr)

    fprs = np.array(fprs)
    tprs = np.array(tprs)
    avg_fpr = fprs.mean(axis=0)
    avg_tpr = tprs.mean(axis=0)
    plt.title('ROC')
    plt.ylabel('True Positive Rate')
    plt.xlabel('False Positive Rate')
    plt.plot(avg_fpr, avg_tpr, 'b')
    plt.show()


sk_predictor()
